#include "refine/window_pass.h"
#include "refine/copy.h"
#include <limits>
#include <iostream>
#include "imglib/imageclass.h"
#include <imglib/bitmap.h>
#include <imglib/bmpmorph.h>
#include <glib/g3/smatrix.h>

using std::clog;
using std::endl;
using std::flush;

namespace refine {
namespace smooth_despill {

typedef unsigned uint;
typedef g3::basic_vector<uint>		vectoru;
typedef g3::basic_vector<float>		vectorf;
typedef g3::basic_smatrix<float>	smatrixf;
typedef g3::basic_smatrix<uint>		smatrixu;

typedef unsigned char uchar;
const uint uchar_max = std::numeric_limits<uchar>::max();
const uint uchar_avg = uchar_max / 2;

template <typename Color>
inline vectoru v(Color const & c)
	{ return vectoru(c.b, c.g, c.r); }

inline uchar limit(float v) {
	return (v < 0) ? 0 : (v >= uchar_max) ? uchar_max : uchar(v + 0.5f);
}

template <typename Color>
inline void discrete(vectorf const & v, Color & c) {
	c.b = limit(v.x);
	c.g = limit(v.y);
	c.r = limit(v.z);
}

// alpha_smoother

class alpha_smoother_base {
public:
	alpha_smoother_base(img::bitmap & modifiable_, float max_condition_number_/*, float more_contrast*/) :
		modifiable(modifiable_),
		max_condition_number(max_condition_number_),
		condition_mult(1 / (max_condition_number - 1))
		//contrast_half(more_contrast / 2),
		//contrast_mult(uchar_max / (1 - more_contrast))
	{ 
	}
	size_extent dim() const
		{ return modifiable.dim(); }
	uint check_count() const
		{ return checks; }
	uint change_count() const
		{ return achanges; }
	bool selected(size_point const & p) const
		{ return modifiable[p]; }
protected:
	/// finishes classification when t0, t1 become known
	unsigned classify_end(float condition_number, float t0, float t1) const
	{
		assert (condition_number <= max_condition_number);
		if (t0 * t1 <= 0) {
			float a = fabs(t0 / (t1 - t0));
			// contrast is increasing with moving away from the bound
			float more_contrast = condition_mult * (condition_number - 1);
			float contrast_half = more_contrast / 2;
			if (a <= contrast_half)
				return 0;
			else if (a >= 1 - contrast_half)
				return uchar_max;
			else
				return (uchar_max / (1 - more_contrast)) * (a - contrast_half) + 0.5f;
		}
		else if (t0 == t1)
			return uchar_avg;
		else if (t0*t0 > t1*t1)
			return uchar_max;
		else
			return 0;
	}
protected:
	uint checks, achanges;
	img::bitmap & modifiable;
	float const max_condition_number;
private:
	float const condition_mult;
//	float const contrast_half, contrast_mult;
};

// despiller

template <typename Raster>
class despiller_base {
public:
	typedef typename Raster::elem_type elem_type;
public:
	despiller_base(
		Raster const & im_, 
		Raster & next_im_,  ///< it receives updates in color channels
		float max_condition_number_)
		: im(im_), next_im(next_im_),
		  max_condition_number(max_condition_number_)
		{ }
	size_extent dim() const
		{ return im.dim(); }
	bool selected(size_point const & p) const {
		uint a = im[p].a;
		return a < uchar_max //because points with maximum alpha will not change in either case
			&& a > 0; //because points with minimum alpha will not be seen
	}
protected:
	Raster const & im;
	Raster & next_im;
	float const max_condition_number;
};

// summators

namespace a {

class summator
{
public:
	summator()
		{ clear(); }
	void clear() {
		N = A = AA = 0;
		I = AI = vectoru::nill;
	}
	template <typename Color>
	summator & operator += (Color const & c) {
		vectoru i = v(c);
		uint a = c.a;
		++N;
		A += a;
		AA += a * a;
		I += i;
		AI += a * i;
		return * this;
	}
	template <typename Color>
	summator & operator -= (Color const & c) {
		vectoru i = v(c);
		uint a = c.a;
		--N;
		A -= a;
		AA -= a * a;
		I -= i;
		AI -= a * i;
		return * this;
	}
	uint count() const
		{ return N; }
	/// calculates mean colors both of foreground and background,
	/// returns false if condition number is too bad for that calculation, and does not return condition number itself in that case
	bool means(float max_condition_number, vectorf & mu1, vectorf & mu0, float * condition_number = 0) const
	{
		//         [ a b ]
		// A^T A = [ b c ]
		float a = AA;
		float b = uchar_max * A - AA;
		float c = uchar_max * uchar_max * N + AA - 2 * uchar_max * A;

		float discriminant = sqr(a - c) + 4 * b * b;
		if (sqr(max_condition_number + 1) * discriminant > sqr((max_condition_number - 1) * (a + c)))
			return false;

		if (condition_number) {
			float sqrt_D = sqrt(discriminant);
			*condition_number = (a + c + sqrt_D) / (a + c - sqrt_D);
		}

		float rdet = uchar_max / (a * c - b * b);
		assert (rdet > 0);

		vectorf sum_I = float(uchar_max) * vectorf(I);
		vectorf sum_alpha_I = vectorf(AI);

		mu1 = rdet * ((c + b) * sum_alpha_I - b * sum_I);
		mu0 = rdet * (-(a + b) * sum_alpha_I + a * sum_I);

		return true;
	}
protected:
	uint N, A, AA; ///< respectively: the count of pixels, sum of alphas, sum of squared alphas
	vectoru I, AI; ///< respectively: sum of colors, sum of (alpha*color)
};

/// these despiller fits both for equal spheres and equal ellipses 
template <typename Raster>
class despiller : public despiller_base<Raster>
{
public:
	despiller(
		Raster const & im_,	Raster & next_im_, 
		float max_condition_number)
		: despiller_base<Raster>(im_, next_im_, max_condition_number)
	{
		clear(); 
	}
	void clear() {
		remove_all();
	}
	void add(size_point const & p)
		{ s += im[p]; }
	void remove(size_point const & p)
		{ s -= im[p]; }
	void remove_all() 
		{ s.clear(); }
	void process(size_point const & p, size_rect const &) 
	{
		vectorf F, B;
		if (!s.means(max_condition_number, F, B)) {
			//next_im[p].a = (im[p].a <= uchar_avg) ? 0 : uchar_max;
			return;
		}

		vectorf I = v(im[p]);
		vectorf d0 = I - B;
		vectorf d1 = I - F;

		uint a = im[p].a; //255 * alpha
		uint b = uchar_max - a; //255 * (1 - alpha)
		float rden = 1.0f / (a * a + b * b);
		vectorf c1 = I + rden * (a * b) * d0 - rden * (b * b) * d1;
		discrete(c1, next_im[p]);
	}
private:
	summator s;
};

} // namespace a

namespace b {

class summator : public a::summator
{
	typedef a::summator base;
public:
	summator()
		{ clear(); }
	void clear() {
		base::clear();
		II.zero();
	}
	template <typename Color>
	summator & operator += (Color const & c) {
		base::operator +=(c);
		II += tsquare(v(c));
		return * this;
	}
	template <typename Color>
	summator & operator -= (Color const & c) {
		base::operator -=(c);
		II -= tsquare(v(c));
		return * this;
	}
	/// calculates mean colors both of foreground and background as well as inversed common covariance matrix,
	/// returns false if condition number is too bad for that calculation
	bool distributions(float max_condition_number, vectorf & mu1, vectorf & mu0, smatrixf & iCov, float * condition_number = 0) const
	{
		if (!base::means(max_condition_number, mu1, mu0, condition_number))
			return false;

		vectorf dmu = mu1 - mu0;

		// (n-1) * uchar_max * Sigma
		smatrixf Sigma = 
			  float(uchar_max * uchar_max) * smatrixf(II) 
			- float(uchar_max * uchar_max * N) * tsquare(mu0)
			- float(uchar_max * A) * sproduct(mu0, dmu)
			- float(AA) * tsquare(dmu);

		float tr = Sigma.tr();
		// really trace can't be negative, since the matrix is nonnegative definite,
		// although it can happen due to round-off errors for null matrix;
		// if the matrix is null then alpha classification is perfect already
		if (tr <= 0) 
			return false;
		Sigma += 0.001f * tr;

		iCov = Sigma.inverse();

		return true;
	}
protected:
	smatrixu II; ///< sum of i*i^T
};

} // namespace b

} //namespace smooth_despill
} //namespace refine
